"""
@File    :  scapy_function.py
@Author  :  LeeSen
@Date    :  2024-08-05
@Version :  Python 3.7.4
"""
import binascii
from scapy.compat import bytes_encode
from scapy.layers.dhcp import DHCPTypes, DHCP, BOOTP
from scapy.layers.inet import IP, ICMP, UDP
from scapy.layers.l2 import Ether, ARP
from scapy.sendrecv import sendp
from my_dhcp_client import DhcpPackageType


def my_send_pkt(pkt, nic):
    sendp(pkt, iface=nic)


def generate_arp_packet(dst, src, hwsrc, psrc, hwdst, pdst, op_code):
    etherHeader = Ether(src=src, dst=dst)
    arp = ARP(hwsrc=hwsrc, psrc=psrc, hwdst=hwdst, pdst=pdst, op=op_code)
    arpPkt = etherHeader / arp
    return arpPkt


def send_icmp_request(src_ip, src_mac, target_ip, target_mac, nic):
    """发送ICMP请求"""
    etherHeader = Ether(src=src_mac, dst=target_mac)
    IPHeader = IP(src=src_ip, dst=target_ip)
    icmpPkt = etherHeader / IPHeader / ICMP()
    sendp(icmpPkt, iface=nic)


def send_dhcp_pkt(dhcp_type_id, dhcp_client, nic):
    bootp_chaddr = binascii.unhexlify(dhcp_client.mac.replace(':', ''))
    dhcp_options = None
    if dhcp_type_id == DhcpPackageType.DISCOVER.value:
        dhcp_options = DHCP(options=[("message-type", DHCPTypes[dhcp_type_id]),
                                     ("client_id", b'\x01' + bootp_chaddr),
                                     ('hostname', "LeeSen"),
                                     ('vendor_class_id', b'Kyland'),
                                     ('param_req_list', [
                                         1, 3, 6, 15, 31, 33, 43, 44, 46, 47, 119, 121, 249, 252
                                     ]),
                                     "end"])
    elif dhcp_type_id == DhcpPackageType.REQUEST.value:
        if dhcp_client.discover_stage:
            dhcp_options = DHCP(options=[("message-type", DHCPTypes[dhcp_type_id]),
                                         ("client_id", b'\x01' + bootp_chaddr),
                                         ("requested_addr", dhcp_client.request_ip),
                                         ("server_id", dhcp_client.dhcp_server_id),
                                         ('hostname', "LeeSen"),
                                         ("client_FQDN", b'\x00\x00\x00' + bytes_encode("LeeSen")),
                                         ('vendor_class_id', b'Kyland'),
                                         ('param_req_list', [
                                             1, 3, 6, 15, 31, 33, 43, 44, 46, 47, 119, 121, 249, 252
                                         ]),
                                         "end"])
        else:
            dhcp_options = DHCP(options=[("message-type", DHCPTypes[dhcp_type_id]),
                                         ("client_id", b'\x01' + bootp_chaddr),
                                         ("server_id", dhcp_client.dhcp_server_id),
                                         ('hostname', "LeeSen"),
                                         ("client_FQDN", b'\x00\x00\x00' + bytes_encode("LeeSen")),
                                         ('vendor_class_id', b'Kyland'),
                                         ('param_req_list', [
                                             1, 3, 6, 15, 31, 33, 43, 44, 46, 47, 119, 121, 249, 252
                                         ]),
                                         "end"])



    elif dhcp_type_id == DhcpPackageType.DECLINE.value:
        dhcp_options = DHCP(options=[("message-type", DHCPTypes[dhcp_type_id]),
                                     ("requested_addr", dhcp_client.request_ip),
                                     ("server_id", dhcp_client.dhcp_server_id),
                                     "end"])
    elif dhcp_type_id == DhcpPackageType.RELEASE.value:
        dhcp_options = DHCP(options=[("message-type", DHCPTypes[dhcp_type_id]),
                                     ("client_id", b'\x01' + bootp_chaddr),
                                     ("server_id", dhcp_client.dhcp_server_id),
                                     "end"])
    bootp_head = None
    if dhcp_type_id == DhcpPackageType.DISCOVER.value or dhcp_type_id == DhcpPackageType.REQUEST.value:
        bootp_head = BOOTP(chaddr=bootp_chaddr, ciaddr=dhcp_client.ip, xid=dhcp_client.transaction_id)
    elif dhcp_type_id == DhcpPackageType.RELEASE.value or dhcp_type_id == DhcpPackageType.DECLINE.value:
        bootp_head = BOOTP(yiaddr=dhcp_client.ip, chaddr=bootp_chaddr, xid=dhcp_client.transaction_id)

    udp_head = UDP(sport=68, dport=67)
    ip_head = IP(src=dhcp_client.ip, dst=dhcp_client.serve_ip)
    ether_head = Ether(src=dhcp_client.mac, dst=dhcp_client.server_mac)
    dhcp_pkt = ether_head / ip_head / udp_head / bootp_head / dhcp_options
    sendp(dhcp_pkt, iface=nic)
    pass


def send_arp(dhcp_client, nic):
    arp_head = ARP(hwsrc=dhcp_client.mac, psrc=dhcp_client.ip, pdst=dhcp_client.ip, op=1)
    ether_head = Ether(src=dhcp_client.mac, dst="ff:ff:ff:ff:ff:ff")
    arp_pkt = ether_head / arp_head
    sendp(arp_pkt, iface=nic)


def send_arp_backup(arp_type, dst, src, hwsrc, psrc, hwdst, pdst, nic):
    arp_head = ARP(hwsrc=hwsrc, psrc=psrc, hwdst=hwdst, pdst=pdst, op=arp_type)
    ether_head = Ether(src=src, dst=dst)
    arp_pkt = ether_head / arp_head
    sendp(arp_pkt, iface=nic)
